package fem2.strategies;

import fem2.Model;
import fem2.element.StructuralElement;
import fem2.enu.EchoLevelType;
import fem2.jlinalg.Solver;
import inf.jlinalg.Array2DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.JArpackSBMGESolver;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public abstract class BucklingStrategy extends Strategy {

	private int bucklingMode;
	protected int nBucklingMode;
	protected JArpackSBMGESolver geSolver;

	public BucklingStrategy(Model m) {
		super(m);
		this.setNumberOfTimeStep(1);
		nBucklingMode = 1;
	}

	@Override
	public double[] getCurrentUHat() {
		IMatrix eiv = (IMatrix) geSolver.getEigenvectors();
		double[] v = new double[eiv.getRowCount()];
		for (int i = 0; i < v.length; i++) {
			v[i] = eiv.get(i, bucklingMode);
		}
		return v;
	}

	/**
	 * get the current eigen buckling mode
	 * 
	 * @return
	 */
	public double getCurrentEigenMode() {
		return geSolver.getEigenvalues().get(bucklingMode);
	}

	/**
	 * set number of buckling modes
	 * 
	 * @param nb
	 */
	public void setNumberOfBucklingModes(int nb) {
		nBucklingMode = nb;
	}

	/**
	 * get number of buckling modes
	 * 
	 * @param nb
	 */
	public int getNumberOfBucklingModes() {
		return nBucklingMode;
	}

	@Override
	public double[] getCurrentUdHat() {
		throw new Error("this method is not available for this strategy");
	}

	@Override
	public double[] getCurrentUddHat() {
		throw new Error("this method is not available for this strategy");
	}

	/**
	 * set the buckling mode to retrieve by getCurrentUHat
	 * 
	 * @param bucklingMode
	 */
	public void setCurrentTimeStep(int bucklingMode) {
		this.bucklingMode = bucklingMode;
	}

	@Override
	public double getCurrentTimeStep() {
		return bucklingMode;
	}

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		super.started(leftHandSideSolver, rightHandSide, u);
		this.geSolver = new JArpackSBMGESolver();
		int n = model.getSize();
		// initialize geA and geB solver
		geSolver.getAInfo().setSize(n);
		// geSolver.getAInfo().setBandwidth(2 *n - 1);
		geSolver.getBInfo().setSize(n);
		// geBInfo.setBandwidth(2 * n- 1);
		// number of eigenvalues and vectors to compute
		geSolver.setNev(nBucklingMode);
		geSolver.setSigma(1);
		geSolver.setMode(JArpackSBMGESolver.MODE_BUCKLING);
		geSolver.initialize();
	}

	@Override
	public void timeStepStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
	}

	@Override
	public void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {

	}

	/**
	 * assemble geometric stiffness matrix of the model
	 * 
	 * @param m
	 * @param kg
	 */
	protected void assembleKg(Model m, IMatrix kg) {
		long start = System.currentTimeMillis();

		/*
		 * get the array of dimension of element stiffness matrix
		 */
		Set<Integer> dimSet = new HashSet<Integer>();
		for (int i = 0; i < m.countElements(); i++) {
			dimSet.add(m.getElement(i).getMeshPart().countDofs());
		}
		/*
		 * create a map from dimension to empty matrix
		 */
		Map<Integer, IMatrix> mapKte = new HashMap<Integer, IMatrix>();
		for (int d : dimSet) {
			IMatrix dummy = new Array2DMatrix(d, d);
			mapKte.put(d, dummy);
		}
		/*
		 * assembly
		 */
		for (int i = 0; i < m.countElements(); i++) {
			StructuralElement e = (StructuralElement) m.getElement(i);
			if (!e.isInactive()) {
				IMatrix kge = mapKte.get(e.getMeshPart().countDofs());
				BLAM.zero(kge);
				kge = (IMatrix) e.computeKge(kge);

				int[] dofs = e.getMeshPart().getDofNumbers();

				for (int j = 0; j < kge.getRowCount(); j++) {
					for (int k = 0; k < kge.getColumnCount(); k++) {
						if (dofs[j] != -1 && dofs[k] != -1) {
							kg.add(dofs[j], dofs[k], kge.get(j, k));
						}
					}
				}
			}
		}

		long end = System.currentTimeMillis();
		if (getEchoLevel().contain(EchoLevelType.AUX1)) {
			System.out.println("Assemble Kg time: " + (end - start) + " ms.");
		}
	}
}
